package cn.cloud.all.security.crypto.encrypt;

import cn.cloud.all.security.crypto.keygen.BytesKeyGenerator;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.modes.AEADBlockCipher;
import org.bouncycastle.crypto.modes.GCMBlockCipher;
import org.bouncycastle.crypto.params.AEADParameters;

import static cn.cloud.all.security.crypto.EncodingUtils.concatenate;
import static cn.cloud.all.security.crypto.EncodingUtils.subArray;

public class BouncyCastleAesGcmBytesEncryptor extends BouncyCastleAesBytesEncryptor {

    public BouncyCastleAesGcmBytesEncryptor(String password, CharSequence salt) {
        super(password, salt);
    }

    public BouncyCastleAesGcmBytesEncryptor(String password, CharSequence salt, BytesKeyGenerator ivGenerator) {
        super(password, salt, ivGenerator);
    }

    @Override
    public byte[] encrypt(byte[] bytes) {
        byte[] iv = this.ivGenerator.generateKey();

        @SuppressWarnings("deprecation")
        GCMBlockCipher blockCipher = new GCMBlockCipher(new org.bouncycastle.crypto.engines.AESFastEngine());
        blockCipher.init(true, new AEADParameters(secretKey, 128, iv, null));

        byte[] encrypted = process(blockCipher, bytes);
        return iv != null ? concatenate(iv, encrypted) : encrypted;
    }

    @Override
    public byte[] decrypt(byte[] encryptedBytes) {
        byte[] iv = subArray(encryptedBytes, 0, this.ivGenerator.getKeyLength());
        encryptedBytes = subArray(encryptedBytes, this.ivGenerator.getKeyLength(), encryptedBytes.length);

        @SuppressWarnings("deprecation")
        GCMBlockCipher blockCipher = new GCMBlockCipher(new org.bouncycastle.crypto.engines.AESFastEngine());
        blockCipher.init(false, new AEADParameters(secretKey, 128, iv, null));
        return process(blockCipher, encryptedBytes);
    }

    private byte[] process(AEADBlockCipher blockCipher, byte[] in) {
        byte[] buf = new byte[blockCipher.getOutputSize(in.length)];
        int bytesWritten = blockCipher.processBytes(in, 0, in.length, buf, 0);
        try {
            bytesWritten += blockCipher.doFinal(buf, bytesWritten);
        } catch (InvalidCipherTextException e) {
            throw new IllegalStateException("unable to encrypt/decrypt", e);
        }
        if (bytesWritten == buf.length) {
            return buf;
        }
        byte[] out = new byte[bytesWritten];
        System.arraycopy(buf, 0, out, 0, bytesWritten);
        return out;
    }
}
